import math
import random
import matplotlib.pyplot as plt
from typing import List, Tuple, Union

# 设置中文字体为黑体，用来正常显示中文标签，并设置能正常显示负号
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# 加载Iris数据集
def load_iris_data() -> List[Tuple[List[float], str]]:
    data = []
    try:
        with open('iris.data', 'r') as f:
            for line in f:
                if line.strip():
                    row = line.strip().split(',')
                    features = list(map(float, row[:-1]))
                    label = row[-1]
                    data.append((features, label))
    except FileNotFoundError:
        print("数据集文件 'iris.data' 不存在，请检查文件路径！")
    return data


# 计算欧氏距离
def euclidean_distance(point1: List[float], point2: List[float]) -> float:
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(point1, point2)))


# KNN分类算法
def knn_classify(train_data: List[Tuple[List[float], str]], test_point: List[float], k: int) -> str:
    distances = []
    for features, label in train_data:
        distance = euclidean_distance(features, test_point)
        distances.append((distance, label))
    distances.sort(key=lambda x: x[0])  # 按距离排序
    nearest_neighbors = distances[:k]
    labels = [label for _, label in nearest_neighbors]
    return max(set(labels), key=labels.count)  # 多数表决


# 划分训练集和测试集
def train_test_split(data: List[Tuple[List[float], str]], test_size: float = 0.2) -> Tuple[
    List[Tuple[List[float], str]], List[Tuple[List[float], str]]]:
    random.shuffle(data)
    split_index = int(len(data) * (1 - test_size))
    return data[:split_index], data[split_index:]


# 可视化分类结果
def visualize_results(test_data: List[Tuple[List[float], str]], predictions: List[str], k: int) -> None:
    colors = {'Iris-setosa': 'r', 'Iris-versicolor': 'g', 'Iris-virginica': 'b'}
    markers = {'correct': 'o', 'wrong': 's'}
    plt.figure(figsize=(10, 6))
    for (features, true_label), pred_label in zip(test_data, predictions):
        is_correct = true_label == pred_label
        plt.scatter(features[0], features[1], c=colors[true_label],
                    marker=markers['correct' if is_correct else 'wrong'])
    plt.title(f'KNN分类结果 (k={k})')
    plt.xlabel('特征 1')
    plt.ylabel('特征 2')
    plt.legend(loc='upper right', labels=['Iris-setosa', 'Iris-versicolor', 'Iris-virginica'])
    plt.show()


# 主函数
if __name__ == "__main__":
    # 加载数据
    data = load_iris_data()
    if not data:
        exit(1)
    train_data, test_data = train_test_split(data)

    k = 3
    predictions = []
    correct = 0
    total = len(test_data)
    class_correct_counts = {'Iris-setosa': 0, 'Iris-versicolor': 0, 'Iris-virginica': 0}
    class_total_counts = {'Iris-setosa': 0, 'Iris-versicolor': 0, 'Iris-virginica': 0}

    for features, label in test_data:
        prediction = knn_classify(train_data, features, k)
        predictions.append(prediction)
        if prediction == label:
            correct += 1
            class_correct_counts[label] += 1
        class_total_counts[label] += 1

    accuracy = correct / total
    print(f"KNN分类准确率: {accuracy * 100:.2f}%")
    print("各类别分类情况：")
    for class_name in class_correct_counts:
        precision = class_correct_counts[class_name] / class_total_counts[class_name] if class_total_counts[
                                                                                             class_name] > 0 else 0
        print(f"{class_name}类别准确率: {precision * 100:.2f}%")

    # 可视化分类结果
    visualize_results(test_data, predictions, k)
